import torch
from torch import nn


class YOLOv1Net(nn.Module):
    def __init__(self, num_classes):
        super(YOLOv1Net, self).__init__()

        self.num_classes = num_classes + 1

        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 64, 7, 2, 3),
            nn.Conv2d(64,64,3,2,1),
            nn.BatchNorm2d(64)
        )

        self.conv2 = nn.Sequential(
            nn.Conv2d(64, 192, 3, 1, 1),
            nn.Conv2d(192, 192, 3, 2, 1),
            nn.BatchNorm2d(192)
        )

        self.conv3 = nn.Sequential(
            nn.Conv2d(192, 128, 1, 1),
            nn.Conv2d(128, 256, 3, 1, 1),
            nn.Conv2d(256, 256, 1, 1),
            nn.Conv2d(256, 512, 3, 1, 1),
            nn.Conv2d(512, 512, 3, 2, 1),
            nn.BatchNorm2d(512)
        )

        self.conv4 = nn.Sequential(
            nn.Conv2d(512, 256, 1, 1),
            nn.Conv2d(256, 512, 3, 1, 1),
            nn.Conv2d(512, 256, 1, 1),
            nn.Conv2d(256, 512, 3, 1, 1),
            nn.Conv2d(512, 256, 1, 1),
            nn.Conv2d(256, 512, 3, 1, 1),
            nn.Conv2d(512, 256, 1, 1),
            nn.Conv2d(256, 512, 3, 1, 1),
            nn.Conv2d(512, 512, 1, 1),
            nn.Conv2d(512, 1024, 3, 1, 1),
            nn.Conv2d(1024, 1024, 3, 2, 1),
            nn.BatchNorm2d(1024)
        )

        self.conv5 = nn.Sequential(
            nn.Conv2d(1024, 512, 1, 1),
            nn.Conv2d(512, 1024, 3, 1, 1),
            nn.Conv2d(1024, 512, 1, 1),
            nn.Conv2d(512, 1024, 3, 1, 1),
            nn.Conv2d(1024, 1024, 3, 2, 1),
            nn.BatchNorm2d(1024)
        )

        self.conv6 = nn.Sequential(
            nn.Conv2d(1024, 1024, 3, 1, 1),
            nn.Conv2d(1024, 1024, 3, 1, 1),
            nn.BatchNorm2d(1024)
        )

        self.conv_fc = nn.Conv2d(1024, 10 + num_classes + 1, 1, 1)
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.relu(self.conv2(x))
        x = self.relu(self.conv3(x))
        x = self.relu(self.conv4(x))
        x = self.relu(self.conv5(x))
        x = self.relu(self.conv6(x))
        x = self.sigmoid(self.conv_fc(x))

        return x.permute([0, 2, 3, 1])


if __name__ == "__main__":
    device = torch.device("cpu")

    model = YOLOv1Net(2)
    model = model.to(device)
    image = torch.randn(1, 3, 448, 448)
    result = model(image.to(device))

    print(result.shape)
